package main

import (
	"context"
	"flag"
	"fmt"
	"net/http"
	"os"
	"path/filepath"
	"registry-proxy/cert"
	"runtime"
	"sync"
	"time"

	"github.com/gorilla/mux"
	"github.com/robfig/cron/v3"
	"github.com/sirupsen/logrus"
)

var (
	port     = flag.String("port", "", "Port to listen on")
	upstream = flag.String("upstream", "", "Upstream registry URL")
	verbose  = flag.Bool("verbose", false, "Enable verbose logging")
)

// 全局变量用于管理服务器实例
var (
	httpServer     *http.Server
	httpsServer    *http.Server
	globalConfig   *Config
	globalRouter   *mux.Router
	serverMutex    sync.RWMutex
	serversRunning bool
)

// stopServer 停止HTTP和HTTPS服务器
func stopServer() {
	serverMutex.Lock()
	defer serverMutex.Unlock()

	if !serversRunning {
		logrus.Info("Servers are already stopped")
		return
	}

	logrus.Info("Stopping servers for certificate update...")

	// 创建超时上下文
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	// 停止HTTPS服务器
	if httpsServer != nil {
		logrus.Info("Stopping HTTPS server...")
		if err := httpsServer.Shutdown(ctx); err != nil {
			logrus.Errorf("Failed to shutdown HTTPS server: %v", err)
		} else {
			logrus.Info("HTTPS server stopped successfully")
		}
		httpsServer = nil
	}

	// 停止HTTP服务器
	if httpServer != nil {
		logrus.Info("Stopping HTTP server...")
		if err := httpServer.Shutdown(ctx); err != nil {
			logrus.Errorf("Failed to shutdown HTTP server: %v", err)
		} else {
			logrus.Info("HTTP server stopped successfully")
		}
		httpServer = nil
	}

	serversRunning = false
	logrus.Info("All servers stopped")
}

// startServer 启动HTTP和HTTPS服务器
func startServer() {
	serverMutex.Lock()
	defer serverMutex.Unlock()

	if serversRunning {
		logrus.Info("Servers are already running")
		return
	}

	if globalConfig == nil || globalRouter == nil {
		logrus.Error("Global config or router is not initialized")
		return
	}

	logrus.Info("Starting servers...")

	// 重新加载证书文件路径（可能已更新）
	httpsEnable := getEnvOrDefault("PROXY_HTTPS_ENABLE", "false")
	if httpsEnable == "true" {
		// 从/cert/目录下查找*.crt和*.key文件并加载证书
		certFiles, err := filepath.Glob("cert/*.crt")
		if err != nil {
			logrus.Errorf("Failed to find certificate files: %v", err)
		} else {
			keyFiles, err := filepath.Glob("cert/*.key")
			if err != nil {
				logrus.Errorf("Failed to find key files: %v", err)
			} else if len(certFiles) > 0 && len(keyFiles) > 0 {
				globalConfig.Server.TlsCert = certFiles[0]
				globalConfig.Server.TlsKey = keyFiles[0]
				logrus.Infof("Updated TLS cert: %s", globalConfig.Server.TlsCert)
				logrus.Infof("Updated TLS key: %s", globalConfig.Server.TlsKey)
			}
		}
	}

	// 启动HTTPS服务器（如果配置了TLS）
	if globalConfig.Server.TlsCert != "" && globalConfig.Server.TlsKey != "" {
		httpsServer = &http.Server{
			Addr:         ":" + globalConfig.Server.HttpSPort,
			Handler:      globalRouter,
			ReadTimeout:  globalConfig.Server.ReadTimeout,
			WriteTimeout: globalConfig.Server.WriteTimeout,
			IdleTimeout:  globalConfig.Server.IdleTimeout,
		}
		logrus.Infof("HTTPS Server starting on port: %s", globalConfig.Server.HttpSPort)
		logrus.Infof("Using TLS cert: %s", globalConfig.Server.TlsCert)
		logrus.Infof("Using TLS key: %s", globalConfig.Server.TlsKey)
		go func() {
			if err := httpsServer.ListenAndServeTLS(globalConfig.Server.TlsCert, globalConfig.Server.TlsKey); err != nil && err != http.ErrServerClosed {
				logrus.Errorf("HTTPS Server failed: %v", err)
			}
		}()
	}

	// 启动HTTP服务器
	httpServer = &http.Server{
		Addr:         ":" + globalConfig.Server.Port,
		Handler:      globalRouter,
		ReadTimeout:  globalConfig.Server.ReadTimeout,
		WriteTimeout: globalConfig.Server.WriteTimeout,
		IdleTimeout:  globalConfig.Server.IdleTimeout,
	}
	logrus.Infof("HTTP Server starting on port: %s", globalConfig.Server.Port)
	go func() {
		if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			logrus.Errorf("HTTP Server failed: %v", err)
		}
	}()

	serversRunning = true
	logrus.Info("All servers started successfully")
}

func main() {
	flag.Parse()

	// Generate default config if requested
	// Load configuration
	config := DefaultConfig()
	config.Logging.Level = "error"
	// Override config with command line flags if provided
	if *port != "" {
		config.Server.Port = *port
	}
	if *upstream != "" {
		config.Registry.Upstream = *upstream
	}
	if *verbose {
		config.Logging.Level = "debug"
	}
	logLevel := os.Getenv("LOG_LEVEL")
	if logLevel != "" {
		config.Logging.Level = logLevel
	}
	enableAutoCert := os.Getenv("ENABLE_AUTO_CERT")
	if enableAutoCert == "1" {
		certDomain := os.Getenv("CERT_DOMAIN")
		certEmail := os.Getenv("CERT_EMAIL")
		if certDomain != "" && certEmail != "" {
			err := cert.InitCertMagic(certEmail, "cert/")
			if err != nil {
				logrus.Fatalf("Failed to initialize certificate magic: %v", err)
			} else {
				err := cert.ObtainCertificate("cert/", certDomain)
				if err != nil {
					logrus.Fatalf("Failed to obtain certificate: %v", err)
				} else {
					//创建定时任务，每天凌晨2点重新获取证书
					task := cron.New()
					task.AddFunc("0 0 2 * * *", func() {
						certPath, keyPath, err := cert.GetCertPaths("cert/", certDomain)
						if err != nil {
							logrus.Errorf("Failed to get cert paths: %v", err)
							return
						} else {
							//检查证书是否过期
							if cert.NeedsRenewal(certPath, keyPath) {
								//停止http ,https服务
								stopServer()

								logrus.Infof("Certificate for domain %s is expired", certDomain)
								//重新获取证书
								err := cert.ObtainCertificate("cert/", certDomain)
								if err != nil {
									logrus.Errorf("Failed to obtain certificate: %v", err)
								} else {
									logrus.Info("Certificate renewed successfully")
								}
								//重新启动http ,https服务
								startServer()
							} else {
								logrus.Info("Certificate is still valid, no renewal needed")
							}
						}
					})
					task.Start()
				}
			}
		}
	}

	// Override with environment variables
	config.Server.Port = getEnvOrDefault("PROXY_PORT", config.Server.Port)
	config.Registry.Upstream = getEnvOrDefault("PROXY_UPSTREAM", config.Registry.Upstream)
	config.Logging.Level = getEnvOrDefault("PROXY_LOG_LEVEL", config.Logging.Level)

	// TLS configuration from environment variables
	httpsEnable := getEnvOrDefault("PROXY_HTTPS_ENABLE", "false")
	if httpsEnable == "true" {
		//从/cert/目录下查找*.crt和*.key文件并加载证书
		certFiles, err := filepath.Glob("cert/*.crt")
		if err != nil {
			logrus.Fatalf("Failed to find certificate files: %v", err)
		}
		keyFiles, err := filepath.Glob("cert/*.key")
		if err != nil {
			logrus.Fatalf("Failed to find key files: %v", err)
		}
		if len(certFiles) == 0 || len(keyFiles) == 0 {
			logrus.Fatalf("Failed to find certificate files: %v", err)
		}
		config.Server.TlsCert = certFiles[0]
		config.Server.TlsKey = keyFiles[0]
		config.Server.Port = "80"
		config.Server.HttpSPort = "443"
		logrus.Infof("Using TLS cert: %s", config.Server.TlsCert)
		logrus.Infof("Using TLS key: %s", config.Server.TlsKey)
	}

	// Force debug logging for detailed proxy analysis
	logrus.Infof("[MAIN] Log level set to Debug for detailed proxy analysis")

	// Validate configuration
	if err := config.Validate(); err != nil {
		logrus.Fatalf("Invalid configuration: %v", err)
	}

	// Setup logging
	setupLogging(config.Logging)

	logrus.WithField("force_output", true).Info("Starting Docker Registry Proxy Server")
	logrus.WithField("force_output", true).Info("Listening on port: ", config.Server.Port)
	logrus.WithField("force_output", true).Info("Listening on HTTPS port: ", config.Server.HttpSPort)
	logrus.WithField("force_output", true).Info("Upstream registry: ", config.Registry.Upstream)

	proxy := NewRegistryProxy(config)

	r := mux.NewRouter()

	// Apply middleware
	r.Use(RecoveryMiddleware)
	r.Use(LoggingMiddleware)
	r.Use(CORSMiddleware)

	// Health check endpoint
	r.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, `{"status":"ok","version":"1.0.0","upstream":"%s"}`, config.Registry.Upstream)
	})

	// Metrics endpoint (basic)
	r.HandleFunc("/metrics", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "text/plain")
		w.WriteHeader(http.StatusOK)
		fmt.Fprintf(w, "# Registry Proxy Metrics\n# TODO: Implement actual metrics\n")
	})

	// Docker Registry API v2 routes
	r.PathPrefix("/v2/").HandlerFunc(proxy.HandleRequest)

	// Handle all other requests (including /token for authentication)
	r.PathPrefix("/").HandlerFunc(proxy.HandleGenericRequest)

	// 设置全局变量
	globalConfig = config
	globalRouter = r

	// 启动服务器
	startServer()

	// 阻塞主程序，等待服务器运行
	select {}
}

func setupLogging(config LoggingConfig) {
	// Set log level
	switch config.Level {
	case "debug":
		logrus.SetLevel(logrus.DebugLevel)
	case "info":
		logrus.SetLevel(logrus.InfoLevel)
	case "warn":
		logrus.SetLevel(logrus.WarnLevel)
	case "error":
		logrus.SetLevel(logrus.ErrorLevel)
	default:
		logrus.SetLevel(logrus.InfoLevel)
	}

	// Set log format
	if config.Format == "json" {
		logrus.SetFormatter(&logrus.JSONFormatter{})
	} else {
		logrus.SetFormatter(&logrus.TextFormatter{
			FullTimestamp: true,
		})
	}
	//设置日志输出格式 包含时间戳2025/08/20 09:05:44 样式 包含调用代码行号
	logrus.SetFormatter(&logrus.TextFormatter{
		FullTimestamp:   true,
		TimestampFormat: "2006/01/02 15:04:05",
		CallerPrettyfier: func(f *runtime.Frame) (string, string) {
			return "", fmt.Sprintf("%s:%d", f.File, f.Line)
		},
	})
}

// getEnvOrDefault returns the value of the environment variable or the default value if not set
func getEnvOrDefault(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}
